import base24;
import config;
import enigma;
import fastrng;
import strutils;
import squeue;
import SimpleThreadingListener;

import core.exception;
import core.sync.mutex;
import core.thread;

import std.algorithm.comparison;
import std.concurrency;
import std.container.slist;
import std.experimental.logger.core;
import std.file;
import std.random;
import std.socket;
import std.typecons;
import std.uni;

import botan.hash.rmd128;
import botan.hash.skein_512;
import botan.libstate.init;
import botan.mac.hmac;
import botan.pubkey.pkcs8;
import botan.pubkey.x509_key;
import botan.rng.rng;



alias EnigmaSetting = Tuple!(string, "rotors", string, "rings", string, "reflector");

const string ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";

const uint HMAC_BLOCK_SIZE = 64;
const uint HMAC_OUTPUT_BITS = 256;

const uint RSA_DATA_LENGTH = 256;

const size_t CONTROL_MESSAGE_PADDING_LENGTH_MAX = 640;

const size_t EDGE_CONNECTIONS_INITIAL_LENGTH = 32;

const size_t HANDSHAKE_INBUF_SIZE = 4096;
const size_t SERVER_INBUF_SIZE = 16384;
const size_t EDGE_INBUF_SIZE = 768;
const size_t CLIENT_INBUF_SIZE = 32768;
const size_t SERVER_OUTBOUND_MESSAGE_SIZE = 1024;

const Duration SERVER_READ_INTERVAL = dur!("msecs")(1);



RandomNumberGenerator cryptoRng;

__gshared PrivateKey privkey;

shared MonoTime[string] recent_pow;
__gshared Mutex recent_pow_mutex;

__gshared Socket client_socket;
__gshared FastXorshiftStar client_fastrng;
__gshared SharedQueue!string client_send_queue;
__gshared SharedQueue!string clean_connections;
__gshared SharedQueue!string[string] ob_queues;

__gshared bool client_shutdown = false;

__gshared EnigmaSetting enigma_setting;



FastXorshiftStar create_fastrng()
{
	char[] s = new char[16];
	for (size_t i = 0; i < 16; i++)
	{
		s[i] = cast(char)uniform(0, 256);
	}
	return FastXorshiftStar(s);
}

char[] fast_random_data(size_t length, FastXorshiftStar rng)
{
	char[] buf = new char[length];
	for (size_t i = 0; i < length; i++)
	{
		buf[i] = rng.get();
	}
	return buf;
}

char[] fast_random_length_random_data(size_t upper_bound, FastXorshiftStar rng)
{
	return fast_random_data(uniform(0, upper_bound), rng);
}

char[] generate_random_rotors()
{
	return [ALPHABET[uniform(0, 26)], ALPHABET[uniform(0, 26)], ALPHABET[uniform(0, 26)]];
}

ubyte[] generate_hmac_key()
{
	ubyte[HMAC_BLOCK_SIZE] key;
	cryptoRng.randomize(key.ptr, HMAC_BLOCK_SIZE);
	return key.dup;
}

ubyte[] random_data(size_t length)
{
	ubyte[] data = new ubyte[length];
	cryptoRng.randomize(data.ptr, length);
	return data;
}

ubyte[] random_length_of_random_data(size_t upper_bound)
{
	return random_data(uniform(0, upper_bound));
}

ubyte[] generate_alphabet()
{
	ubyte[] _in = cast(ubyte[])(ALPHABET.dup);
	randomShuffle(_in);
	return _in;
}

ubyte[] random_letters(size_t length)
{
	ubyte[] buf = new ubyte[length];
	for (size_t i = 0; i < length; i++)
	{
		buf[i] = cast(ubyte)uniform(65, 91);
	}
	return buf;
}

ubyte[] random_length_of_random_letters(size_t upper_bound)
{
	return random_letters(uniform(0, upper_bound));
}

string add_random_letters(string src, string pool)
{
	char[] buf = new char[src.length + 30]; // assume 30 chars will be added
	size_t i = 0;
	size_t j = 0;
	size_t u = pool.length;
	size_t rand, ol, sl;
	while (i < src.length)
	{
		rand = uniform(15, 25); // so hard encoding

		sl = i + rand;
		if (sl > src.length)
		{
			rand = src.length - i;
			buf.length = j + rand;
			buf[j .. (j + rand)] = src[i .. (i + rand)];
			return assumeUnique(buf);
		}

		ol = j + rand + 1;

		if (ol > buf.length)
		{
			buf.length = ol + 16; // so magic number
		}
		buf[j .. ol] = src[i .. sl] ~ pool[uniform(0, u)];
		i = sl;
		j = ol;
	}
	buf.length = j;
	return assumeUnique(buf);
}

string add_spaces(string src)
{
	char[] buf = new char[src.length + 152]; // assume 152 spaces will be added
	size_t i = 0;
	size_t j = 0;
	size_t rand, ol, sl;
	while (i < src.length)
	{
		rand = uniform(1, 25); // still hard encoding

		sl = i + rand;
		if (sl > src.length)
		{
			rand = src.length - i;
			buf.length = j + rand;
			buf[j .. (j + rand)] = src[i .. (i + rand)];
			return assumeUnique(buf);
		}

		ol = j + rand + 1;

		if (ol > buf.length)
		{
			buf.length = ol + 12; // yet another magic number
		}
		buf[j .. ol] = src[i .. sl] ~ ' ';
		i = sl;
		j = ol;
	}
	buf.length = j;
	return assumeUnique(buf);
}

string random_cased(string src) // input should be uppercase
{
	char[] buf = cast(char[])unsafe_lower(src);
	for (size_t i = uniform(0, 10); i < src.length; i += uniform(0, 10))
	{
		byte b = cast(ubyte)buf[i];
		if (b >= 97 && b < 123)
		{
			buf[i] = cast(char)(b - 32);
		}
	}
	return assumeUnique(buf);
}



bool checkPOW(ref SecureVector!ubyte hash, int leadingZero, int maxNB)
{
	if (leadingZero >= hash.length)
	{
		return false;
	}
	for (int i = 0; i < leadingZero; i++)
	{
		if (hash[i] != 0)
		{
			return false;
		}
	}
	return hash[leadingZero] <= maxNB;
}



void requestConnection()
{
	client_send_queue.put(assumeUnique('c' ~
			fast_random_length_random_data(CONTROL_MESSAGE_PADDING_LENGTH_MAX, client_fastrng)));
}

void clientSenderThread(string rotors, immutable ubyte[] hmac_key, string base24_alphabet)
{
	auto enigma = EnigmaMachine.fromKeySheet(enigma_setting.expand);
	enigma.setDisplay(rotors);

	auto hmac = new HMAC(new Skein512(HMAC_OUTPUT_BITS));
	hmac.setKey(hmac_key.ptr, HMAC_BLOCK_SIZE);

	long len;

	while (true)
	{
		if (client_shutdown)
		{
			client_socket.close();
			return;
		}

		string msg = client_send_queue.get();

		if (msg is null)
		{
			client_socket.close();
			return;
		}

		msg = cast(string)pack_uint(cast(uint)msg.length) ~ msg;

		hmac.update(msg);
		SecureVector!ubyte mac = hmac.finished();

		msg = enigma.processText(enbase(mac[] ~ cast(ubyte[])msg, base24_alphabet));

		msg = add_spaces(random_cased(msg));

		len = client_socket.send(msg);

		if (len == client_socket.ERROR || len < msg.length)
		{
			trace("Client shutdown, client_socket send error");
			client_socket.close();
			client_shutdown = true;
			return;
		}
	}
}

void clientReceiverThread(string rotors, immutable ubyte[] hmac_key, string base24_alphabet)
{
	auto enigma = EnigmaMachine.fromKeySheet(enigma_setting.expand);
	enigma.setDisplay(rotors);

	auto hmac = new HMAC(new Skein512(HMAC_OUTPUT_BITS));
	hmac.setKey(hmac_key.ptr, HMAC_BLOCK_SIZE);

	char[] buf = new char[CLIENT_INBUF_SIZE];
	long len;

	char unmatched = '\0';
	string seg;
	ubyte[] fin_mac = new ubyte[HMAC_OUTPUT_BITS/8];

	while (true)
	{
		if (client_shutdown)
		{
			client_socket.close();
			return;
		}

		len = client_socket.receive(buf);

		if (len == 0 || len == client_socket.ERROR)
		{
			trace("Client shutdown, client_socket recv error");
			client_socket.close();
			client_shutdown = true;
			return;
		}

		string msg = remove_spaces(buf[0 .. len], len);
		string text;

		try
		{
			msg = toUpper(msg);

			if (unmatched)
			{
				text = unmatched ~ enigma.processText(msg);
				unmatched = '\0';
			}
			else
			{
				text = enigma.processText(msg);
			}

			if (text.length % 2 != 0)
			{
				unmatched = text[$ - 1];
				text = text[0 .. $ - 1];
			}

			seg ~= debase(text, base24_alphabet);
		}
		catch
		{
			trace("Client shutdown, error decoding message");
			client_socket.shutdown(SocketShutdown.BOTH);
			client_socket.close();
			client_shutdown = true;
			return;
		}

		while (true)
		{
			if (seg.length < 37)
			{
				break;
			}

			uint body_len = unpack_uint(cast(ubyte[])seg[32 .. 36]);

			if (seg.length < body_len + 36)
			{
				break;
			}

			ubyte[] mac = cast(ubyte[])seg[0 .. 32];
			string content = seg[32 .. 36 + body_len];
			seg = seg[36 + body_len .. $];

			hmac.update(content);
			hmac.flushInto(fin_mac);

			if (!equal(fin_mac, mac))
			{
				trace("Client shutdown, bad hmac");
				client_socket.shutdown(SocketShutdown.BOTH);
				client_socket.close();
				client_shutdown = true;
				return;
			}

			try
			{
				final switch(content[4])
				{
					case 'c':
						if (content.length < 9)
						{
							trace("Client shutdown, incorrect message length");
							client_socket.shutdown(SocketShutdown.BOTH);
							client_socket.close();
							client_shutdown = true;
							return;
						}

						ob_queues[content[5 .. 9]] = new SharedQueue!string();
						clean_connections.put(content[5 .. 9]);
						break;
					case 'd':
						if (content.length < 9)
						{
							trace("Client shutdown, incorrect message length");
							client_socket.shutdown(SocketShutdown.BOTH);
							client_socket.close();
							client_shutdown = true;
							return;
						}

						auto ptr = content[5 .. 9] in ob_queues;

						if (ptr is null)
						{
							trace("Client shutdown, incorrect cid");
							client_socket.shutdown(SocketShutdown.BOTH);
							client_socket.close();
							client_shutdown = true;
							return;
						}

						(*ptr).destroy();
						break;
					case 'm':
						if (content.length <= 9)
						{
							trace("Client shutdown, incorrect message length");
							client_socket.shutdown(SocketShutdown.BOTH);
							client_socket.close();
							client_shutdown = true;
							return;
						}

						auto ptr = content[5 .. 9] in ob_queues;

						if (ptr is null)
						{
							trace("Client shutdown, incorrect cid");
							client_socket.shutdown(SocketShutdown.BOTH);
							client_socket.close();
							client_shutdown = true;
							return;
						}

						(*ptr).put(content[9 .. $]);
						break;
				}
			}
			catch (SwitchError)
			{
				trace("Client shutdown, unknown cmd");
				client_socket.shutdown(SocketShutdown.BOTH);
				client_socket.close();
				client_shutdown = true;
				return;
			}
		}
	}
}

void clientConnectionSender(shared Socket socket, string cid)
{
	Socket sock = cast(Socket)socket;

	long len;

	while (true)
	{
		if (client_shutdown)
		{
			sock.close();
			return;
		}

		auto data = ob_queues[cid].get();

		if (data is null)
		{
			sock.close();
			return;
		}

		len = sock.send(data);

		if (len == sock.ERROR || len < data.length)
		{
			sock.close();
			client_send_queue.put(assumeUnique('d' ~ cid ~
					fast_random_length_random_data(CONTROL_MESSAGE_PADDING_LENGTH_MAX, client_fastrng)));
			return;
		}
	}
}

void clientConnectionHandler(shared Socket socket)
{
	Socket sock = cast(Socket)socket;

	trace("Connection from ", sock.remoteAddress);

	requestConnection();

	string cid = clean_connections.get(dur!("seconds")(config.getUnsignedShort("timeout")));

	if (cid is null)
	{
		sock.close();
		return;
	}

	trace("Got connection id " ~ to!string(unpack_uint(cast(ubyte[])cid)));

	Thread.sleep(SERVER_READ_INTERVAL * 2);

	spawn(&clientConnectionSender, socket, cid);

	char[] buf = new char[EDGE_INBUF_SIZE];

	long len;

	while (true)
	{
		if (client_shutdown)
		{
			sock.close();
			return;
		}

		len = sock.receive(buf);

		if (len == sock.ERROR || len == 0)
		{
			sock.close();
			client_send_queue.put(assumeUnique('d' ~ cid ~
					fast_random_length_random_data(CONTROL_MESSAGE_PADDING_LENGTH_MAX, client_fastrng)));
			return;
		}

		client_send_queue.put(assumeUnique('m' ~ cid ~ buf[0 .. len]));
	}
}


void launchClient(PublicKey pubkey)
{
	info("Checking RSA public key...");

	if (!pubkey.checkKey(cryptoRng, false))
	{
		warning("Invaild key, encryption may be broken");
	}

	info("Doing Proof of Work...");

	string pow;
	SecureVector!ubyte pow_hash;

	int leading_zeroes = config.getUnsignedShort("powleadingzero");
	int maxNB = config.getUnsignedShort("powfirstbytemax");
	string pow_salt = config.getString("powsalt");
	if (pow_salt is null)
	{
		pow_salt = "";
	}

	do
	{
		RIPEMD128 pow_hasher = new RIPEMD128;
		pow = random_cased(cast(string)random_letters(32));
		pow_hasher.update(pow);
		pow_hasher.update(pow_salt);
		pow_hash = pow_hasher.finished();
	} while (!checkPOW(pow_hash, leading_zeroes, maxNB));

	pow_hash.clear();

	info("Connecting...");

	try
	{
		auto timeout = dur!"seconds"(config.getUnsignedShort("timeout"));
		client_socket = new Socket(AddressFamily.INET, SocketType.STREAM);
		client_socket.setOption(SocketOptionLevel.SOCKET, SocketOption.RCVTIMEO, timeout);
		client_socket.setOption(SocketOptionLevel.SOCKET, SocketOption.SNDTIMEO, timeout);
		client_socket.connect(new InternetAddress(config.getString("remote"), config.getUnsignedShort("rport")));
	}
	catch (SocketException)
	{
		fatal("Could not connect to server");
	}

	handshake(pubkey, pow);
}

void handshake(PublicKey pubkey, string pow)
{
	char[] buf = new char[HANDSHAKE_INBUF_SIZE];
	long len;

	string first_msg = add_spaces(pow ~ random_cased(cast(string)random_length_of_random_letters(640)));

	len = client_socket.send(first_msg);

	if (len == client_socket.ERROR)
	{
		client_socket.close();
		fatal("Could not write on socket, " ~ client_socket.getErrorText());
	}

	if (len < first_msg.length)
	{
		client_socket.shutdown(SocketShutdown.BOTH);
		client_socket.close();
		fatal("Could not dispose a message immediately");
	}

	info("Handshaking...");

	len = client_socket.receive(buf);

	if (len == 0)
	{
		client_socket.close();
		fatal("Connection closed before cookie received");
	}

	if (len == client_socket.ERROR)
	{
		client_socket.close();
		fatal("Could not read socket, " ~ client_socket.getErrorText());
	}

	if (len < 34)
	{
		client_socket.shutdown(SocketShutdown.BOTH);
		client_socket.close();
		fatal("Cookie too short");
	}

	string r_msg = remove_spaces(buf, 34);

	if (r_msg.length < 34)
	{
		client_socket.shutdown(SocketShutdown.BOTH);
		client_socket.close();
		fatal("Cookie too short");
	}

	string cookie = toUpper(r_msg[0 .. 8]);
	char[] rotors = generate_random_rotors();
	ubyte[] hmac_key = generate_hmac_key();

	string alphabet = toUpper(r_msg[8 .. 34]);
	string base24_alphabet = alphabet[0 .. 24];
	string unused_letters = alphabet[24 .. 26];

	ubyte[] keys = cast(ubyte[])cookie ~ cast(ubyte[])rotors ~ hmac_key;

	PKEncryptorEME encryptor = new PKEncryptorEME(pubkey, "EME1(SHA-256)");

	ubyte[] encrypted_keys = encryptor.encrypt(keys.ptr, keys.length, cryptoRng)[] ~ random_length_of_random_data(224);

	delete encryptor;

	string msg = add_spaces(random_cased(add_random_letters(enbase(encrypted_keys, base24_alphabet), unused_letters)));

	len = client_socket.send(msg);

	if (len == client_socket.ERROR)
	{
		client_socket.close();
		fatal("Could not write on socket, " ~ client_socket.getErrorText());
	}

	if (len < msg.length)
	{
		client_socket.shutdown(SocketShutdown.BOTH);
		client_socket.close();
		fatal("Could not dispose a message immediately");
	}

	client_socket.setOption(SocketOptionLevel.SOCKET, SocketOption.RCVTIMEO, dur!"seconds"(0));

	info("Connection established");

	delete buf;

	clientMain(rotors, hmac_key, base24_alphabet);
}

void clientMain(char[] rotors, ubyte[] hmac_key, string base24_alphabet)
{
	client_fastrng = create_fastrng();

	client_send_queue = new SharedQueue!string();
	clean_connections = new SharedQueue!string();

	spawn(&clientSenderThread, rotors.idup, hmac_key.idup, base24_alphabet.idup);
	spawn(&clientReceiverThread, rotors.idup, hmac_key.idup, base24_alphabet.idup);

	auto listener = new SimpleThreadingListener(config.getString("listen"), config.getUnsignedShort("port"), &clientConnectionHandler);
	listener.start();

	info("Listening on " ~ config.getString("listen") ~ ":" ~ to!string(config.getUnsignedShort("port")));

	while (true)
	{
		if (client_shutdown)
		{
			client_send_queue.destroy();
			clean_connections.destroy();
			foreach (q; ob_queues)
			{
				q.destroy();
			}
			fatal("Disconnected");
		}
		Thread.sleep(dur!("seconds")(1));
	}
}



void serverConnectionHandler(shared Socket socket)
{
	Socket sock = cast(Socket)socket; // the compiler doesn't know it is unique

	trace("Connection from ", sock.remoteAddress);

	char[] buf = new char[HANDSHAKE_INBUF_SIZE];
	long len;

	Duration timeout_duration = dur!"seconds"(config.getUnsignedShort("timeout"));

	sock.setOption(SocketOptionLevel.SOCKET, SocketOption.SNDTIMEO, dur!"seconds"(2));
	sock.setOption(SocketOptionLevel.SOCKET, SocketOption.RCVTIMEO, timeout_duration);

	len = sock.receive(buf);

	if (len == 0 || len == sock.ERROR)
	{
		sock.close();
		return;
	}

	if (len < 32)
	{
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}

	string pow = remove_spaces(buf, 32);

	auto now = MonoTime.currTime();
	auto pow_lifetime = dur!"minutes"(config.getUnsignedShort("powlife"));

	synchronized (recent_pow_mutex)
	{
		auto ptr = pow in recent_pow;
		if (ptr !is null)
		{
			info("Client connecting with PoW seen before");
			*ptr = now;
			sock.shutdown(SocketShutdown.BOTH);
			sock.close();
			return;
		}
	}

	string pow_salt = config.getString("powsalt");
	if (pow_salt is null)
	{
		pow_salt = "";
	}

	RIPEMD128 pow_hasher = new RIPEMD128;
	pow_hasher.update(pow);
	pow_hasher.update(pow_salt);
	auto pow_hash = pow_hasher.finished();
	if (!checkPOW(pow_hash, config.getUnsignedShort("powleadingzero"), config.getUnsignedShort("powfirstbytemax")))
	{
		info("Client connecting with bad PoW");
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}

	synchronized (recent_pow_mutex)
	{
		recent_pow[pow] = MonoTime.currTime();

		foreach (pow_time; recent_pow.byKeyValue())
		{
			if (now - pow_time.value > pow_lifetime)
			{
				recent_pow.remove(pow_time.key);
			}
		}
	}

	delete pow_hasher;
	pow_hash.clear();

	ubyte[] cookie = random_letters(8);
	ubyte[] alphabet = generate_alphabet();

	string base24_alphabet = cast(string)alphabet[0 .. 24];
	string unused_letters = cast(string)alphabet[24 .. 26];

	string data = add_spaces(random_cased(cast(string)(cookie ~ alphabet ~
			random_length_of_random_letters(CONTROL_MESSAGE_PADDING_LENGTH_MAX))));

	len = sock.send(data);

	if (len == sock.ERROR)
	{
		sock.close();
		return;
	}

	if (len < data.length)
	{
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}

	delete data;

	len = sock.receive(buf);

	if (len == 0 || len == sock.ERROR)
	{
		sock.close();
		return;
	}

	if (len < RSA_DATA_LENGTH * 2)
	{
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}

	char[] raw_msg = strip_all(buf, " " ~ unused_letters ~ unsafe_lower(unused_letters), RSA_DATA_LENGTH * 2);

	toUpperInPlace(raw_msg);

	if (raw_msg.length < RSA_DATA_LENGTH * 2)
	{
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}

	ubyte[] rsa_msg;
	SecureVector!ubyte decrypted_msg;
	PKDecryptorEME decryptor;
	try
	{
		decryptor = new PKDecryptorEME(privkey, "EME1(SHA-256)");
		rsa_msg = cast(ubyte[])debase(raw_msg[0 .. RSA_DATA_LENGTH * 2], base24_alphabet);
		decrypted_msg = decryptor.decrypt(rsa_msg.ptr, rsa_msg.length);
	}
	catch
	{
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}
	finally
	{
		delete decryptor;
	}


	if (!equal(decrypted_msg[0 .. 8], cookie))
	{
		sock.shutdown(SocketShutdown.BOTH);
		sock.close();
		return;
	}
	char[] rotors = cast(char[])decrypted_msg[8 .. 11];
	ubyte[] hmac_key = decrypted_msg[11 .. 11 + HMAC_BLOCK_SIZE];

	delete cookie;
	delete raw_msg;
	delete rsa_msg;
	decrypted_msg.clear();

	auto upstream_enigma = EnigmaMachine.fromKeySheet(enigma_setting.expand); // send
	auto downstream_enigma = EnigmaMachine.fromKeySheet(enigma_setting.expand); // recv
	upstream_enigma.setDisplay(rotors);
	downstream_enigma.setDisplay(rotors);

	auto upstream_hmac = new HMAC(new Skein512(HMAC_OUTPUT_BITS)); // send
	auto downstream_hmac = new HMAC(new Skein512(HMAC_OUTPUT_BITS)); // recv
	upstream_hmac.setKey(hmac_key.ptr, HMAC_BLOCK_SIZE);
	downstream_hmac.setKey(hmac_key.ptr, HMAC_BLOCK_SIZE);

	scope (exit)
	{
		delete upstream_hmac;
		delete downstream_hmac;
	}

	FastXorshiftStar fastrng = create_fastrng();

	sock.setOption(SocketOptionLevel.SOCKET, SocketOption.SNDTIMEO, dur!"seconds"(0));
	sock.setOption(SocketOptionLevel.SOCKET, SocketOption.RCVTIMEO, dur!"seconds"(0));
	sock.blocking(false);

	string remote_addr = config.getString("remote");
	ushort remote_port = config.getUnsignedShort("rport");

	Socket[] edge_connections;
	MonoTime[] edge_connection_initial;
	edge_connections.length = EDGE_CONNECTIONS_INITIAL_LENGTH;
	edge_connection_initial.length = EDGE_CONNECTIONS_INITIAL_LENGTH;
	uint curr_cid = 0;

	ushort idle_timeout = config.getUnsignedShort("idletimeout");

	scope (exit)
	{
		foreach (conn; edge_connections)
		{
			if (conn !is null)
			{
				conn.close();
			}
		}
	}

	buf = new char[SERVER_INBUF_SIZE];
	char[] inbuf = new char[EDGE_INBUF_SIZE];
	char[] ib_msg;
	ubyte[] fin_mac = new ubyte[HMAC_OUTPUT_BITS/8];
	auto send_queue = SList!string();
	string text, raw, seg, content, outbuf;
	char unmatched = '\0';
	SocketSet read_set = new SocketSet();
	SocketSet write_set = new SocketSet();
	SocketSet except_set = new SocketSet(); // write is not used
	while (true)
	{
		read_set.reset();
		read_set.add(sock);
		except_set.reset();
		except_set.add(sock);
		foreach (conn; edge_connections)
		{
			if (conn !is null)
			{
				read_set.add(conn);
				except_set.add(conn);
			}
		}

		if (idle_timeout == 0)
		{
			len = Socket.select(read_set, write_set, except_set);
		}
		else
		{
			len = Socket.select(read_set, write_set, except_set, dur!"minutes"(idle_timeout));
		}


		if (len == -1)
		{
			continue;
		}

		if (read_set.isSet(sock) || except_set.isSet(sock))
		{

			len = sock.receive(buf);

			if (len == 0 || len == sock.ERROR)
			{
				sock.close();
				return;
			}
		}
		else
		{
			if (len == 0)
			{
				sock.close();
				return;
			}
			len = 0;
		}

		if (len)
		{
			ib_msg = cast(char[])remove_spaces(buf[0 .. len], len);

			try
			{
				toUpperInPlace(ib_msg);

				if (unmatched)
				{
					text = unmatched ~ downstream_enigma.processText(cast(immutable)(ib_msg));
					unmatched = '\0';
				}
				else
				{
					text = downstream_enigma.processText(cast(immutable)(ib_msg));
				}

				if (text.length % 2 != 0)
				{
					unmatched = text[$ - 1];
					text = text[0 .. $ - 1];
				}

				raw = debase(text, base24_alphabet);
			}
			catch
			{
				trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", error decoding message");
				sock.shutdown(SocketShutdown.BOTH);
				sock.close();
				return;
			}
		}
		else
		{
			raw = "";
		}

		seg ~= raw;

		while (true)
		{
			if (seg.length < 37)
			{
				break;
			}

			uint body_len = unpack_uint(cast(ubyte[])seg[32 .. 36]);

			if (seg.length < body_len + 36)
			{
				break;
			}

			ubyte[] mac = cast(ubyte[])seg[0 .. 32];
			content = seg[32 .. 36 + body_len];
			seg = seg[36 + body_len .. $];

			downstream_hmac.update(content);
			downstream_hmac.flushInto(fin_mac);

			if (!equal(fin_mac, mac))
			{
				trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", bad hmac");
				sock.shutdown(SocketShutdown.BOTH);
				sock.close();
				return;
			}

			try
			{
				final switch(content[4])
				{
					case 'c':
						if (curr_cid == edge_connections.length)
						{
							edge_connections.length += EDGE_CONNECTIONS_INITIAL_LENGTH;
							edge_connection_initial.length += EDGE_CONNECTIONS_INITIAL_LENGTH;
						}

						auto edge_sock = new Socket(AddressFamily.INET, SocketType.STREAM);
						edge_sock.blocking(false);
						edge_sock.connect(new InternetAddress(remote_addr, remote_port));
						edge_connections[curr_cid] = edge_sock;
						edge_connection_initial[curr_cid] = MonoTime.currTime;
						send_queue.insertFront(assumeUnique('c' ~ pack_uint(curr_cid) ~
								fast_random_length_random_data(CONTROL_MESSAGE_PADDING_LENGTH_MAX, fastrng)));
						++curr_cid;
						break;
					case 'd':
						if (content.length < 9)
						{
							trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", incorrect message length");
							sock.shutdown(SocketShutdown.BOTH);
							sock.close();
							return;
						}

						uint cid = unpack_uint(cast(ubyte[])content[5 .. 9]);

						if (cid >= curr_cid)
						{
							trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", incorrect cid");
							sock.shutdown(SocketShutdown.BOTH);
							sock.close();
							return;
						}

						if (edge_connections[cid] is null)
						{
							continue;
						}

						edge_connections[cid].shutdown(SocketShutdown.BOTH);
						edge_connections[cid].close();
						edge_connections[cid] = null;
						break;
					case 'm':
						if (content.length <= 9)
						{
							trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", incorrect message length");
							sock.shutdown(SocketShutdown.BOTH);
							sock.close();
							return;
						}

						uint cid = unpack_uint(cast(ubyte[])content[5 .. 9]);

						if (cid >= curr_cid)
						{
							trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", incorrect cid");
							sock.shutdown(SocketShutdown.BOTH);
							sock.close();
							return;
						}

						if (edge_connections[cid] is null)
						{
							continue;
						}

						if (edge_connections[cid].send(content[9 .. $]) == edge_connections[cid].ERROR && !wouldHaveBlocked() &&
								MonoTime.currTime - edge_connection_initial[cid] >= timeout_duration)
						{
							edge_connections[cid].close();
							edge_connections[cid] = null;
							send_queue.insertFront(assumeUnique('d' ~ content[5 .. 9] ~
									fast_random_length_random_data(CONTROL_MESSAGE_PADDING_LENGTH_MAX, fastrng)));
						}
						// drop if couldn't send everything at once, who cares?
				}
			}
			catch (SwitchError)
			{
				trace("Disconnecting " ~ to!string(sock.remoteAddress) ~ ", unknown cmd");
				sock.shutdown(SocketShutdown.BOTH);
				sock.close();
				return;
			}
		}

		if (!outbuf)
		{
			for (uint i = 0; i < curr_cid; i++)
			{
				if (edge_connections[i] is null)
				{
					continue;
				}

				if (read_set.isSet(edge_connections[i]) || except_set.isSet(edge_connections[i]))
				{

					len = edge_connections[i].receive(inbuf);

					if (len == 0 || len == edge_connections[i].ERROR)
					{
						edge_connections[i].close();
						edge_connections[i] = null;
						send_queue.insertFront(assumeUnique('d' ~ pack_uint(i) ~
								fast_random_length_random_data(CONTROL_MESSAGE_PADDING_LENGTH_MAX, fastrng)));
						continue;
					}

					send_queue.insertFront(assumeUnique('m' ~ pack_uint(i) ~ inbuf[0 .. len]));
				}
			}
		}

		Thread.sleep(SERVER_READ_INTERVAL);

		if (outbuf)
		{
			len = sock.send(outbuf);

			if (len == sock.ERROR)
			{
				if (wouldHaveBlocked())
				{
					continue;
				}
				sock.close();
				return;
			}

			if (len == outbuf.length)
			{
				outbuf = "";
			}
			else
			{
				outbuf = outbuf[0 .. len];
				continue;
			}
		}

		while (!send_queue.empty)
		{
			string msg = send_queue.removeAny(); // It should be random, but it just removes the front currently
			msg = cast(string)pack_uint(cast(uint)msg.length) ~ msg;

			upstream_hmac.update(msg);
			SecureVector!ubyte msg1_mac = upstream_hmac.finished();

			msg = upstream_enigma.processText(enbase(msg1_mac[] ~ cast(ubyte[])msg, base24_alphabet));

			while (msg.length < SERVER_OUTBOUND_MESSAGE_SIZE && !send_queue.empty)
			{
				string msg2 = send_queue.removeAny();
				msg2 = cast(string)pack_uint(cast(uint)msg2.length) ~ msg2;

				upstream_hmac.update(msg2);
				SecureVector!ubyte msg2_mac = upstream_hmac.finished();

				msg ~= upstream_enigma.processText(enbase(msg2_mac[] ~ cast(ubyte[])msg2, base24_alphabet));
			}

			msg = add_spaces(random_cased(msg));

			len = sock.send(msg);

			if (len == sock.ERROR)
			{
				if (wouldHaveBlocked())
				{
					outbuf = msg;
					break;
				}
				sock.close();
				return;
			}

			if (len < msg.length)
			{
				outbuf = msg[0 .. len];
				break;
			}
		}
	}
}

void launchServer()
{
	recent_pow_mutex = new Mutex;

	info("Initializing listener...");

	auto listener = new SimpleThreadingListener(config.getString("listen"), config.getUnsignedShort("port"), &serverConnectionHandler);
	listener.start();

	info("Listening on " ~ config.getString("listen") ~ ":" ~ to!string(config.getUnsignedShort("port")));

	listener.join();
}

void initEnigmaSetting()
{
	enigma_setting.rotors = config.getString("rotors");
	enigma_setting.rings = config.getString("rings");
	enigma_setting.reflector = config.getString("reflector");
}

int main()
{
	string config_file;

	globalLogLevel(LogLevel.info);
	info("Reading flint.config");
	try
	{
		config_file = assumeUnique(cast(char[])read("flint.config", 1024));
	}
	catch (FileException)
	{
		fatal("Could not read flint.config");
	}

	readConfig(config_file);
	delete config_file;
	validateConfig();

	LibraryInitializer botan_init;

	cryptoRng = RandomNumberGenerator.makeRng();

	switch (config.getString("type"))
	{
		case "client":
			initEnigmaSetting();
			info("Reading " ~ config.getString("keyfile"));
			auto pubkey = loadKey(config.getString("keyfile"));
			try
			{
				launchClient(pubkey);
			}
			catch (Error)
			{
				// do nothing
			}
			catch (Exception ex)
			{
				error(ex.toString());
			}
			delete pubkey;
			break;
		case "server":
			initEnigmaSetting();
			info("Reading " ~ config.getString("keyfile"));
			privkey = loadKey(config.getString("keyfile"), cryptoRng);
			launchServer();
			break;
		default:
			fatal("Unknown running type: should be 'client' or 'server'");
	}

	delete privkey;
	delete cryptoRng;
	return 0;
}